#!/usr/bin/env python
"""
:mod:`pynchon.bio.io.pdb`
=========================

Provides read/write support for the pdb file format.
"""
import re
import numpy as np
from itertools import izip
from collections import defaultdict
from tempfile import NamedTemporaryFile
from pynchon.bio.data.alphabet import IUPAC_PROTEIN_ONE
from pynchon.bio.atom import Atom, AtomData, Model, ModelData

PDB_COORDS = "%s%5i %-4s%c%3s %c%4i%c   %8.3f%8.3f%8.3f%6.2f%6.2f          %2s%2s\n"
PDB_TER = "%s%5i %-4s%c%3s %c%4i%c\n"

def _parse_atom_line(line):
    """
    (internal) Parses a PDBv3 line. Returns a tuple of (coords, fields). Fields
    are in the order of the line and are valid for ``AtomData`` instances.
    
    Arguments:
    
        - line (``str``) A "ATOM  " or "HETATM" line.
    """
    at_type = line[0:6]
    ser_num = int(line[6:11])   # numbers are ints's?
    at_id = line[12:16]         # " N B "
    alt_loc = line[16:17]       # always keep \s
    res_name = line[17:20]      # non standard is 4chars long
    chain_id = line[21:22]
    res_id = int(line[22:26].strip())   # pdb requirement int
    res_ic = line[26:27]
    x = float(line[30:38])            
    y = float(line[38:46])            
    z = float(line[46:54])
    occupancy = float(line[54:60])
    bfactor = float(line[60:66])
    element = line[76:78]
    charge = line[78:80]
    coords = (x, y, z)
    fields = (at_type, ser_num, at_id, alt_loc, res_name, chain_id, \
              res_id, res_ic, occupancy, bfactor, element, charge)

    return (coords, fields)

def pdb_header_parser(lines):
    header = {}
    seqres = defaultdict(list)
    for line in lines:
        if line.startswith('SEQRES'):
            chain_id = line[11]
            seqres[chain_id].extend([IUPAC_PROTEIN_ONE[aa3.capitalize()] for \
                                       aa3 in line[19:71].split()])
    header['SEQRES'] = dict(seqres)
    return header

def pdb_parser(fh, coords_vector =None, fields_vector =None):
    """
    Parses a PDB file and yields tuples of (coordinates, fields, metadata).
    By default coordinates are returned as a rank-2 numpy array of 64-bit 
    floats, fields i.e. the data associated with each "ATOM" or "HETATM" line in
    the PDB file is returned as a list of tuples representing the fields of a 
    line, metadata is a dict of headers etc.. If "coords_vector"  or 
    "fields_vector" is specified the given ``Entity`` sub-classes will be used
    as containers for atomic coordinates and fields. Typically these should be
    for example the tuples ``(Model, Atom)`` and ``(ModelData, AtomData)`` are 
    valid "coords_vector" and "fields_vector" arguments.  
    
    Argumets:
    
        - fh (``file``) An open for reading file handle or lines in the PDB file
          format.
        - coords_vector (``tuple``) A pair of ``Entity`` sub-classes to store 
          coordinates.
        - fields_vector (``tuple``) A pair of ``Entity`` sub-classes to store
          atomo fields.
    
    """
    coords_vector = coords_vector or (None, None)
    fields_vector = fields_vector or (None, None)
    
    match_coords = re.compile('^HETATM|ATOM|MODEL')    # start of coordinates
    
    lines = fh.readlines()
    c_offset = None
    for i, line in enumerate(lines):
        if match_coords.match(line):
            c_offset = i
            break
    raw_header = lines[:c_offset] if c_offset else ""
    parsed_header = pdb_header_parser(raw_header)

    current_atoms = []
    current_atoms_meta = []
    current_model_id = 0
    for line in lines[c_offset:]:
        record_type = line[0:6]

        if record_type == 'MODEL ' or \
          (record_type == 'ATOM  ' and not current_model_id):
            current_model_id += 1
            model_meta = {'mod_id':current_model_id,
                          'header':parsed_header,
                          'raw_header':raw_header}
            
        if record_type == 'ATOM  ' or record_type == 'HETATM':
            atom_coords, atom_metadata = _parse_atom_line(line)
            if coords_vector[1]:
                current_atoms.append(coords_vector[1](atom_coords, \
                                                metadata={'pdb':atom_metadata}))
            else:
                current_atoms.append(atom_coords)
            if fields_vector[1]:
                current_atoms_meta.append(fields_vector[1](atom_metadata))
            else:
                current_atoms_meta.append(atom_metadata)
                
        elif record_type in ('ENDMDL', 'END   ') and current_atoms:
            # quick
            if not coords_vector[1]:
                current_atoms = np.array(current_atoms, \
                        dtype=getattr(coords_vector[0], 'data_type', 'float64'))
            if not fields_vector[1] and fields_vector[0]:
                current_atoms_meta = np.array(current_atoms_meta, \
                        dtype=fields_vector[0].data_type)
            # meta
            if fields_vector[0]:
                meta = fields_vector[0](current_atoms_meta, copy=False, metadata=model_meta)
            else:
                meta = current_atoms_meta             
            if coords_vector[0]:
                coords = coords_vector[0](current_atoms, copy=False)
            else:
                coords = current_atoms           
            yield (coords, meta, model_meta)
            current_atoms = []
            current_atoms_meta = [] 
        else:
            continue
                
def pdb_writer(fh, coords_vector, fields_vector, metadata =None):
    """
    Writes a PDB file given the contents of vectors of coordinates and fields 
    and a metadata dictionary . The "coords_vector" should be a rank-2 array of
    64 bit floats or a ``AtomVector`` instance. The "fields_vector" should be
    a sequence of tuples representing the 12 non-coordinate fields in a PDB 
    file, see  ``pynchon.bio.atom.AtomData`` the function ``_parse_atom_line`` 
    and for details. Metadata should be a dictionary it overrides 
    ``fields_vector.metadata``.
    
    Arguments:
    
        - fh (``file``) An open for writing file handle or object with ``write``
          and ``writelines`` methods.
        - coords_vector (``Model`` instance or ``numpy.ndarray`` instance). The 
          numpy array should be rank-2 and of lenght N.
        - fields_vector (``ModelData`` instance or ``sequence``) The sequence 
          should be of length N and contain tuples of the 12 non-coordinate 
          fields of a atom line.
        - metadata (``dict``) [default =``None``] If "metadata" is not given 
          ``fields_vector.metadata`` is used.
    
    """
    old_meta = (None,)* 12
    metadata = metadata or getattr(fields_vector, 'metadata', {})
    fh.writelines(metadata.get('raw_header', ()))
    # write first MODEL
    fh.write('MODEL        %s\n' % metadata.get('mod_id', ' '))
    for coords, meta in izip(coords_vector, fields_vector):
        if old_meta[5] and (old_meta[5] != meta[5]) and \
        (old_meta[4].capitalize() in IUPAC_PROTEIN_ONE):
            fh.write(PDB_TER % (('TER   ',) + (old_meta[1] + 1,) + \
                                ('   ', ' ') + old_meta[4:8]))
        fh.write(PDB_COORDS % (meta[:8] + tuple(coords) + meta[8:]))
        old_meta = meta
    # write last TER
    if meta[4].capitalize() in IUPAC_PROTEIN_ONE:
        fh.write(PDB_TER % (('TER   ',) + (meta[1] + 1,) + \
                            ('   ', ' ') + meta[4:8]))
    # write ENDMDL and last END
    fh.write('ENDMDL\n')      
    fh.write('END   \n')
    
def tmp_pdb(coords_vector, fields_vector, metadata =None, **kwargs):
    """
    Writes a ``AtomVector`` and ``AtomDataVector`` to a temporary PDB file. 
    Returns a file handle open for reading and writing. If the handle is closed 
    or  garbage collected the temporary file will be deleted from the file 
    system.
    
    Arguments
    
        - coords_vector (``AtomVector`` or ``numpy.ndarray``) This should be a 
          valid  input for ``pdb_writer`` "coords_vector.
        - fields_vector (``AtomDataVector`` or ``sequence``) This should be a 
          valid  input for ``pdb_writer`` "fields_vector".
        - metadata (``dict``) This should be a valid  input for ``pdb_writer`` 
          "metadata".
          
    Additional keyworded arguments are passed to ``pdb_writer``.
    """
    tmp_file = NamedTemporaryFile()
    pdb_writer(tmp_file, coords_vector, fields_vector, metadata, **kwargs)
    tmp_file.flush()
    return tmp_file

